﻿Imports System.Security.Cryptography
Imports System.Security.Cryptography.X509Certificates
Imports System.Security
Imports System.Xml
Imports System.Security.Cryptography.Xml
Imports System.Xml.Serialization
Imports System.Xml.Schema
Imports System.Linq
Imports Aricie.Security.Cryptography

Namespace Security
    <Serializable> _
    Public Class SerializableSecureString
        Implements IXmlSerializable
        Implements IDisposable
#Region "Constructors"

        ' Default constructor used only by Serializer
        Private Sub New()
        End Sub


        ''' <summary>
        '''     Initialize new instance 
        ''' </summary>
        ''' <param name="instance"></param>
        Public Sub New(instance As SerializableSecureString, Optional bCopyData As Boolean = False)
            Certificate = instance.Certificate
            CertificateStore = instance.CertificateStore
            ThumbPrint = instance.ThumbPrint
            Content = If(bCopyData, instance.CloneData(), New SecureString())
        End Sub

        ''' <summary>
        '''     Externally supplied certificate constructor
        '''     This is likely the production constructor of choice
        '''     Certificate is installed in specified store if it is not already there
        ''' </summary>
        ''' <param name="cert"></param>
        ''' <param name="store"></param>
        Public Sub New(cert As X509Certificate2, store As X509Store)
            CertificateStore = store
            Certificate = cert
            ThumbPrint = cert.Thumbprint
            SetCertificateInStore(Certificate)
        End Sub

#End Region

#Region "Properties"

        ''' <summary>
        '''     Set true by testing constructor, this causes the IDispose implementation
        '''     to purge the autogenerated constructor and its key from the machine
        ''' </summary>
        Private ReadOnly _autoGenerated As Boolean

        ''' <summary>
        '''     The wrappered SecureString
        ''' </summary>
        Private ReadOnly Content As New SecureString()

        ' Design note: The get accessor is to allow instances to share
        ' the dynamically created certificates in unit tests. Whether or not
        ' it remains in production is uncertain.
        Public Property Certificate() As X509Certificate2
            Get
                Return m_Certificate
            End Get
            Private Set(value As X509Certificate2)
                m_Certificate = value
            End Set
        End Property
        Private m_Certificate As X509Certificate2

        ' Certificate Store in which Certificate resides
        ' Default is CurrentUser/Personnal (MY)
        Public Property CertificateStore() As X509Store
            Get
                Return m_CertificateStore
            End Get
            Private Set(value As X509Store)
                m_CertificateStore = value
            End Set
        End Property
        Private m_CertificateStore As X509Store

        ' During deserialization ThumbPrint is needed acquire the Certificate
        ' from CertificateStore
        ' Otherwise we could use Certificate.ThumbPrint.
        Public Property ThumbPrint() As String
            Get
                Return m_ThumbPrint
            End Get
            Private Set(value As String)
                m_ThumbPrint = value
            End Set
        End Property
        Private m_ThumbPrint As String

        ''' <summary>
        '''     Moves Content into and out of serialization streams
        ''' </summary>
        Private Property Value() As XmlDocument
            ' Move the SecureString into the serialization in encrypted xml
            Get
                Dim doc = New XmlDocument()
                doc.LoadXml(SerializationXsi)

                ' copy Content into document
                'A priori knowledge of location of ClearText element
                Dim elmt As XmlNode = doc.DocumentElement.FirstChild.FirstChild
                elmt.InnerText = Extract()

                Dim xmlEnc = New EncryptedXml(doc)
                Dim encrData As EncryptedData = xmlEnc.Encrypt(TryCast(elmt, XmlElement), Certificate)
                EncryptedXml.ReplaceElement(TryCast(elmt, XmlElement), encrData, False)

                LoadCertificateContext(doc)
                Return doc
            End Get
            Set(value As XmlDocument)
                ' Decrypt the xml and reload the SecureString
                Dim doc As XmlDocument = value
                Dim bReadOnly As Boolean = UnloadCertificateContext(doc)

                ' Decrypt the document
                Dim encrXml = New EncryptedXml(doc)
                encrXml.DecryptDocument()

                ' load the secure string
                Initialize(doc.DocumentElement.FirstChild.FirstChild.InnerText)
                If bReadOnly Then
                    Content.MakeReadOnly()
                End If
            End Set
        End Property

#End Region

#Region "Serialization Helpers"

        'XNnames of elements in deserialization schema
        Private Const _Envelope As String = "SecureStringEnvelope"
        Private Const _ClearText As String = "ClearText"
        Private Const _ThumbPrint As String = "ThumbPrint"
        Private Const _Store As String = "Store"
        Private Const _Location As String = "Location"
        Private Const _Certificate As String = "CertificateContext"
        Private Const _IsReadOnly As String = "IsReadOnly"

        'XPath locations in used in queries
        Private Const xp_descendants As String = "descendant::"
        Private Const xp_child As String = "child::"

        'XPath queries used during deserialization and serialization
        Private Const xp_CertificateContext As String = xp_descendants & _Certificate
        Private Const xp_Store As String = xp_child & _Store
        Private Const xp_Location As String = xp_child & _Location
        Private Const xp_ThumbPrint As String = xp_child & _ThumbPrint
        Private Const xp_IsReadOnly As String = xp_descendants & _IsReadOnly

        ' instance of the serialization schema
        'Public Shared SerializationXsi As String = (Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString((Convert.ToString("<") & _Envelope) + ">" + "<EncryptionContext>" + "<") & _ClearText) + "/>" + "<") & _IsReadOnly) + ">false" + "</") & _IsReadOnly) + ">" + "<") & _Certificate) + ">" + "<") & _ThumbPrint) + "/>" + "<") & _Store) + "/>" + "<") & _Location) + "/>" + "</") & _Certificate) + ">" + "</EncryptionContext>" + "</") & _Envelope) + ">"
        Public Shared SerializationXsi As String = String.Format("<{0}><EncryptionContext><{1}/><{2}>false</{2}><{3}><{4}/><{5}/><{6}/></{3}></EncryptionContext></{0}>" _
                                                                 , _Envelope _
                                                                 , _ClearText _
                                                                 , _IsReadOnly _
                                                                 , _Certificate _
                                                                 , _ThumbPrint _
                                                                 , _Store _
                                                                 , _Location)

        ''' <summary>
        ''' Fill in the remainder of the serialization schema
        ''' </summary>
        ''' <param name="doc">The outgoing serialization document</param>
        Private Sub LoadCertificateContext(doc As XmlDocument)
            Dim nsmgr = New XmlNamespaceManager(doc.NameTable)
            Dim elmt As XmlElement = Nothing

            If Content.IsReadOnly() Then
                elmt = TryCast(doc.DocumentElement.SelectSingleNode(xp_IsReadOnly, nsmgr), XmlElement)
                elmt.InnerText = "true"
            End If
            Dim certCtx As XmlNode = doc.DocumentElement.SelectSingleNode(xp_CertificateContext, nsmgr)

            elmt = TryCast(certCtx.SelectSingleNode(xp_ThumbPrint, nsmgr), XmlElement)
            elmt.InnerText = ThumbPrint
            Certificate = Find(ThumbPrint, CertificateStore).FirstOrDefault()

            elmt = TryCast(certCtx.SelectSingleNode(xp_Store, nsmgr), XmlElement)
            elmt.InnerText = CertificateStore.Name

            elmt = TryCast(certCtx.SelectSingleNode(xp_Location, nsmgr), XmlElement)
            elmt.InnerText = [Enum].Format(GetType(StoreLocation), CertificateStore.Location, "G")
        End Sub

        ''' <summary>
        ''' Repopulate class members from the incoming document
        ''' </summary>
        ''' <param name="doc">the incoming serialization document</param>
        ''' <returns>ReadOnly condition for new Content</returns>
        Private Function UnloadCertificateContext(doc As XmlDocument) As Boolean
            Dim nsmgr = New XmlNamespaceManager(doc.NameTable)
            Dim elmt As XmlElement = Nothing

            elmt = TryCast(doc.DocumentElement.SelectSingleNode(xp_IsReadOnly, nsmgr), XmlElement)
            Dim result As Boolean = Convert.ToBoolean(elmt.InnerText)

            Dim certCtx = TryCast(doc.DocumentElement.SelectSingleNode(xp_CertificateContext, nsmgr), XmlElement)

            elmt = TryCast(certCtx.SelectSingleNode(xp_ThumbPrint, nsmgr), XmlElement)
            ThumbPrint = elmt.InnerText

            elmt = TryCast(certCtx.SelectSingleNode(xp_Store, nsmgr), XmlElement)
            Dim xStore As String = elmt.InnerText

            elmt = TryCast(certCtx.SelectSingleNode(xp_Location, nsmgr), XmlElement)
            Dim xStoreLoc As String = elmt.InnerText

            Dim sName = DirectCast([Enum].Parse(GetType(StoreName), xStore, True), StoreName)
            Dim sLoc = DirectCast([Enum].Parse(GetType(StoreLocation), xStoreLoc, True), StoreLocation)
            CertificateStore = New X509Store(sName, sLoc)

            ' fault occurs (deliberately and eventually) if private key does not exist

            Certificate = Find(ThumbPrint).FirstOrDefault()
            If Certificate Is Nothing Then
                Throw New CertificateException(CertificateException.CertificateNotFoundMessage(ThumbPrint, CertificateStore))
            End If
            If Not Certificate.HasPrivateKey Then
                Throw New CertificateException(CertificateException.NoPrivateKeyMessage(ThumbPrint))
            End If
            Return result
        End Function

#End Region

#Region "Certificate Store helpers"

        Private Sub SetCertificateInStore(certificate As X509Certificate2)
            ' for testing we'll use the CurrentUser's store
            ' in production this is likely to be LocalMachine which would 
            ' require DACL allowing current user to wrie

            ' see if the current cert is already there. Use of thumbprint is 
            ' mandatory since all dynamic certs have same name, but each would
            ' have a unique thumbprint.
            If Not Find(certificate.Thumbprint).Any() Then
                CertificateStore.Open(OpenFlags.ReadWrite)
                CertificateStore.Add(certificate)
                CertificateStore.Close()
            End If
        End Sub

        Public Function Find(thumbPrint As String) As IEnumerable(Of X509Certificate2)
            Return Find(thumbPrint, CertificateStore)
        End Function

        Public Shared Function Find(thumbPrint As String, store As X509Store) As IEnumerable(Of X509Certificate2)
            store.Open(OpenFlags.[ReadOnly])
            Dim results As IEnumerable(Of X509Certificate2) = store.Certificates.Find(X509FindType.FindByThumbprint, thumbPrint, False).Cast(Of X509Certificate2)()
            store.Close()
            Return results
        End Function

#End Region

#Region "Content helpers"

        ''' <summary>
        '''     Clear and load initial data into the SecureString
        ''' </summary>
        ''' <param name="clearText"></param>
        Public Sub Initialize(clearText As String)
            Content.Clear()
            Append(clearText)
        End Sub

        ''' <summary>
        '''     Append to current SecureString content
        '''     Throws ReadOnlyContentException if SecureString has been locked
        ''' </summary>
        ''' <param name="clearText">text to append</param>
        Public Sub Append(clearText As String)
            If clearText IsNot Nothing Then
                If Content.IsReadOnly() Then
                    Throw New ReadOnlyContentException()
                End If
                ' iterate the loop to avoid making a copy of the 
                ' cleartext accidentally
                For Each t As Char In clearText
                    Content.AppendChar(t)
                Next
            End If
        End Sub

        ''' <summary>
        '''     Retrieve the data currently within the SecureString
        '''     Caller has the responsibility for keeping this data hopefully
        '''     in a GC0 collection usage.
        ''' </summary>
        ''' <returns>ClearText from SecureString</returns>
        Public Function Extract() As String
            Return CryptoHelper.ReadSecureString(Me.Content)
        End Function

        ''' <summary>
        '''     Seal the SecureString from further modification
        ''' </summary>
        Public Sub MakeReadOnly()
            Content.MakeReadOnly()
        End Sub

        ''' <summary>
        '''     Check read only state of SecureString
        ''' </summary>
        ''' <returns>false if data can be added</returns>
        Public Function IsReadOnly() As Boolean
            Return Content.IsReadOnly()
        End Function

        ''' <summary>
        '''     Return a copy of the internal SecureString
        ''' </summary>
        ''' <returns></returns>
        Public Function CloneData() As SecureString
            Return Content.Copy()
        End Function

#End Region

#Region "IDisposable and the Dispose pattern"

        Private _disposed As Boolean

        ''' <summary>
        '''     Used for housekeeping the autogenerated certificates in test runs
        ''' </summary>
        Public Sub Dispose() Implements IDisposable.Dispose
            Dispose(True)
            GC.SuppressFinalize(Me)
        End Sub

        Protected Overridable Sub Dispose(isDisposing As Boolean)
            If Not _disposed Then
                If isDisposing Then
                    Content.Dispose()
                    ' kill the protected data now
                    If Certificate IsNot Nothing AndAlso _autoGenerated Then
                        ' house keep the autogenerated certificates from disk
                        ' _autogenerated flag can be true only in DEBUG build
                        If Certificate.HasPrivateKey Then
                            ' delete the private key
                            Dim pvKey = TryCast(Certificate.PrivateKey, RSACryptoServiceProvider)
                            pvKey.PersistKeyInCsp = False
                            pvKey.Clear()
                        End If
                        CertificateStore.Open(OpenFlags.ReadWrite)
                        CertificateStore.Remove(Certificate)
                        CertificateStore.Close()
                    End If
                End If
                _disposed = True
            End If
        End Sub

#End Region

#Region "IXmlSerializable"

        Public Function GetSchema() As XmlSchema Implements IXmlSerializable.GetSchema
            Return Nothing
        End Function



        Public Sub ReadXml(reader As XmlReader) Implements IXmlSerializable.ReadXml
            Dim doc = New XmlDocument()

            ' position reader to the start of our serialization stream
            ' because at entry we are in the DataContractSerializer's envelope
            reader.MoveToContent()
            While reader.Read()
                If reader.NodeType = XmlNodeType.Element AndAlso reader.LocalName = _Envelope Then
                    Exit While
                End If
            End While

            doc.LoadXml(reader.ReadOuterXml())

            Value = doc
        End Sub




        Public Sub WriteXml(writer As XmlWriter) Implements IXmlSerializable.WriteXml
            Value.Save(writer)
        End Sub

#End Region


    End Class




End Namespace


